"""
SQLite数据库模型和初始化
"""
import sqlite3
import os
from datetime import datetime
from typing import Optional, List, Dict
import json
import sys


class Database:
    def __init__(self, db_path: str = "ai_assistant.db"):
        self.db_path = db_path
        # 只在数据库文件不存在时才初始化
        if not os.path.exists(db_path):
            print(f"数据库文件不存在，正在创建: {db_path}")
            self.init_database()
        else:
            # 数据库已存在，只确保表结构存在（CREATE TABLE IF NOT EXISTS 会自动处理）
            self._ensure_tables()

    def get_connection(self):
        """获取数据库连接"""
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row  # 使返回结果为字典格式
        return conn

    def init_database(self):
        """初始化数据库表结构"""
        conn = self.get_connection()
        cursor = conn.cursor()

        # 1. 会话表 - 存储对话会话
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS conversations (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                title TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)

        # 2. 消息表 - 存储对话消息
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS messages (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                conversation_id INTEGER NOT NULL,
                role TEXT NOT NULL,  -- 'user' or 'assistant'
                content TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (conversation_id) REFERENCES conversations(id)
            )
        """)

        # 3. 异步任务表 - 存储异步任务状态（视频生成、文件解析等）
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS tasks (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                task_id TEXT UNIQUE NOT NULL,  -- 智谱API返回的任务ID
                task_type TEXT NOT NULL,  -- 'video_generation', 'file_parser', 'chat_async'
                status TEXT DEFAULT 'pending',  -- 'pending', 'processing', 'completed', 'failed'
                result TEXT,  -- 存储结果JSON
                error_message TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                conversation_id INTEGER,
                FOREIGN KEY (conversation_id) REFERENCES conversations(id)
            )
        """)

        # 5. 生成记录表 - 存储生成的图片/视频记录
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS generations (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                generation_type TEXT NOT NULL,  -- 'image' or 'video'
                prompt TEXT NOT NULL,
                result_url TEXT,  -- 生成的资源URL
                local_path TEXT,  -- 本地保存路径
                status TEXT DEFAULT 'pending',
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                conversation_id INTEGER,
                FOREIGN KEY (conversation_id) REFERENCES conversations(id)
            )
        """)

        # 6. 搜索历史表 - 存储网络搜索记录
        # 已移除：搜索历史表（当前不再提供网络搜索）

        # 7. 语音记录表 - 存储语音转文本/文本转语音的记录
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS audio_records (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                audio_type TEXT NOT NULL,  -- 'transcription' or 'speech'
                text_content TEXT,
                audio_path TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                conversation_id INTEGER,
                FOREIGN KEY (conversation_id) REFERENCES conversations(id)
            )
        """)

        # 8. 设置表 - 存储应用设置
        self._ensure_settings_table(conn)

        # 创建索引以提高查询性能
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_messages_conversation ON messages(conversation_id)")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_tasks_task_id ON tasks(task_id)")
        cursor.execute("CREATE INDEX IF NOT EXISTS idx_tasks_status ON tasks(status)")
        # 已移除：文件状态索引

        conn.commit()
        conn.close()
        print(f"数据库初始化完成: {self.db_path}")
        
        # 确保默认会话（ID=1）存在
        self._ensure_default_conversation()
        # 确保默认设置存在
        self._ensure_default_settings()

    def _ensure_default_conversation(self):
        """确保默认会话（ID=1）存在"""
        conn = self.get_connection()
        cursor = conn.cursor()
        try:
            # 检查会话ID=1是否存在
            cursor.execute("SELECT id FROM conversations WHERE id = 1")
            if cursor.fetchone() is None:
                # 如果不存在，先检查是否还有其他会话
                cursor.execute("SELECT COUNT(*) FROM conversations")
                count = cursor.fetchone()[0]
                
                if count == 0:
                    # 没有任何会话，可以安全地插入ID=1
                    # 需要先更新sqlite_sequence表（如果存在）来允许插入ID=1
                    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='sqlite_sequence'")
                    if cursor.fetchone():
                        # 如果sqlite_sequence表存在，更新它
                        cursor.execute("UPDATE sqlite_sequence SET seq = 0 WHERE name = 'conversations'")
                    # 插入ID=1的会话
                    cursor.execute(
                        "INSERT INTO conversations (id, title) VALUES (1, 'AI对话')"
                    )
                    conn.commit()
                    print("已创建默认会话（ID=1）")
                else:
                    # 已经有其他会话了，无法插入ID=1
                    # 这种情况下，我们需要使用第一个会话的ID
                    cursor.execute("SELECT MIN(id) FROM conversations")
                    min_id = cursor.fetchone()[0]
                    if min_id:
                        print(f"警告：无法创建ID=1的会话，现有最小ID为{min_id}，将使用该会话")
        except Exception as e:
            print(f"确保默认会话时出错: {e}")
            import traceback
            traceback.print_exc()
        finally:
            conn.close()

    def _ensure_tables(self):
        """确保数据库表存在（如果表不存在则创建）"""
        conn = self.get_connection()
        cursor = conn.cursor()
        
        # 使用 CREATE TABLE IF NOT EXISTS 确保表存在
        # 这里只检查关键表，其他表在 init_database 中已经使用 IF NOT EXISTS
        try:
            cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='conversations'")
            if cursor.fetchone() is None:
                # 表不存在，需要初始化
                conn.close()
                self.init_database()
            else:
                # 暂时保持连接，用于检查/创建其他表
                self._ensure_settings_table(conn)
                conn.commit()
                conn.close()
                # 表存在，确保默认会话和设置存在
                self._ensure_default_conversation()
                self._ensure_default_settings()
        except Exception as e:
            conn.close()
            print(f"检查数据库表时出错: {e}")
            # 如果出错，尝试重新初始化
            self.init_database()

    # ============ 会话相关操作 ============
    def create_conversation(self, title: str = "新对话") -> int:
        """创建新会话"""
        conn = self.get_connection()
        cursor = conn.cursor()
        cursor.execute(
            "INSERT INTO conversations (title) VALUES (?)",
            (title,)
        )
        conversation_id = cursor.lastrowid
        conn.commit()
        conn.close()
        return conversation_id

    def get_conversations(self) -> List[Dict]:
        """获取所有会话列表"""
        conn = self.get_connection()
        cursor = conn.cursor()
        cursor.execute("""
            SELECT id, title, created_at, updated_at 
            FROM conversations 
            ORDER BY updated_at DESC
        """)
        rows = cursor.fetchall()
        conn.close()
        return [dict(row) for row in rows]

    def get_conversation(self, conversation_id: int) -> Optional[Dict]:
        """获取单个会话信息"""
        conn = self.get_connection()
        cursor = conn.cursor()
        cursor.execute(
            "SELECT * FROM conversations WHERE id = ?",
            (conversation_id,)
        )
        row = cursor.fetchone()
        conn.close()
        return dict(row) if row else None

    def update_conversation_title(self, conversation_id: int, title: str):
        """更新会话标题"""
        conn = self.get_connection()
        cursor = conn.cursor()
        cursor.execute(
            "UPDATE conversations SET title = ?, updated_at = ? WHERE id = ?",
            (title, datetime.now(), conversation_id)
        )
        conn.commit()
        conn.close()

    def delete_conversation(self, conversation_id: int):
        """删除会话（级联删除相关消息）"""
        conn = self.get_connection()
        cursor = conn.cursor()
        # 先删除相关消息
        cursor.execute("DELETE FROM messages WHERE conversation_id = ?", (conversation_id,))
        # 再删除会话
        cursor.execute("DELETE FROM conversations WHERE id = ?", (conversation_id,))
        conn.commit()
        conn.close()

    def clear_messages(self, conversation_id: int):
        """清除会话的所有消息（保留会话）"""
        conn = self.get_connection()
        cursor = conn.cursor()
        # 只删除消息，不删除会话
        cursor.execute("DELETE FROM messages WHERE conversation_id = ?", (conversation_id,))
        # 更新会话的更新时间
        cursor.execute(
            "UPDATE conversations SET updated_at = ? WHERE id = ?",
            (datetime.now(), conversation_id)
        )
        conn.commit()
        conn.close()

    # ============ 消息相关操作 ============
    def add_message(self, conversation_id: int, role: str, content: str) -> int:
        """添加消息"""
        conn = self.get_connection()
        cursor = conn.cursor()
        cursor.execute(
            "INSERT INTO messages (conversation_id, role, content) VALUES (?, ?, ?)",
            (conversation_id, role, content)
        )
        message_id = cursor.lastrowid
        # 更新会话的更新时间
        cursor.execute(
            "UPDATE conversations SET updated_at = ? WHERE id = ?",
            (datetime.now(), conversation_id)
        )
        conn.commit()
        conn.close()
        return message_id

    def get_messages(self, conversation_id: int) -> List[Dict]:
        """获取会话的所有消息"""
        conn = self.get_connection()
        cursor = conn.cursor()
        cursor.execute("""
            SELECT id, role, content, created_at 
            FROM messages 
            WHERE conversation_id = ? 
            ORDER BY created_at ASC
        """, (conversation_id,))
        rows = cursor.fetchall()
        conn.close()
        return [dict(row) for row in rows]

    # ============ 文件相关操作 ============
    # 已移除：文件相关操作（上传与解析功能已删除）

    # ============ 任务相关操作 ============
    def add_task(self, task_id: str, task_type: str, conversation_id: int = None) -> int:
        """添加异步任务"""
        conn = self.get_connection()
        cursor = conn.cursor()
        cursor.execute("""
            INSERT INTO tasks (task_id, task_type, conversation_id)
            VALUES (?, ?, ?)
        """, (task_id, task_type, conversation_id))
        task_db_id = cursor.lastrowid
        conn.commit()
        conn.close()
        return task_db_id

    def update_task_status(self, task_id: str, status: str, result: str = None, 
                          error_message: str = None):
        """更新任务状态"""
        conn = self.get_connection()
        cursor = conn.cursor()
        cursor.execute("""
            UPDATE tasks 
            SET status = ?, result = ?, error_message = ?, updated_at = ?
            WHERE task_id = ?
        """, (status, result, error_message, datetime.now(), task_id))
        conn.commit()
        conn.close()

    def get_task(self, task_id: str) -> Optional[Dict]:
        """获取任务信息"""
        conn = self.get_connection()
        cursor = conn.cursor()
        cursor.execute("SELECT * FROM tasks WHERE task_id = ?", (task_id,))
        row = cursor.fetchone()
        conn.close()
        return dict(row) if row else None

    # ============ 生成记录相关操作 ============
    def add_generation(self, generation_type: str, prompt: str, 
                      conversation_id: int = None) -> int:
        """添加生成记录"""
        conn = self.get_connection()
        cursor = conn.cursor()
        cursor.execute("""
            INSERT INTO generations (generation_type, prompt, conversation_id)
            VALUES (?, ?, ?)
        """, (generation_type, prompt, conversation_id))
        gen_id = cursor.lastrowid
        conn.commit()
        conn.close()
        return gen_id

    def update_generation(self, gen_id: int, result_url: str = None, 
                         local_path: str = None, status: str = None):
        """更新生成记录"""
        conn = self.get_connection()
        cursor = conn.cursor()
        updates = []
        params = []
        if result_url:
            updates.append("result_url = ?")
            params.append(result_url)
        if local_path:
            updates.append("local_path = ?")
            params.append(local_path)
        if status:
            updates.append("status = ?")
            params.append(status)
        params.append(gen_id)
        cursor.execute(
            f"UPDATE generations SET {', '.join(updates)} WHERE id = ?",
            params
        )
        conn.commit()
        conn.close()

    # ============ 搜索历史相关操作 ============
    # 已移除：搜索历史相关操作

    # ============ 语音记录相关操作 ============
    def add_audio_record(self, audio_type: str, text_content: str = None, 
                        audio_path: str = None, conversation_id: int = None) -> int:
        """添加语音记录"""
        conn = self.get_connection()
        cursor = conn.cursor()
        cursor.execute("""
            INSERT INTO audio_records (audio_type, text_content, audio_path, conversation_id)
            VALUES (?, ?, ?, ?)
        """, (audio_type, text_content, audio_path, conversation_id))
        audio_id = cursor.lastrowid
        conn.commit()
        conn.close()
        return audio_id

    # ============ 设置相关操作 ============
    # ============ 设置相关操作 ============
    def get_setting(self, setting_key: str) -> Optional[Dict]:
        """获取设置"""
        conn = self.get_connection()
        cursor = conn.cursor()
        cursor.execute("SELECT setting_key, setting_value, updated_at FROM settings WHERE setting_key = ?", (setting_key,))
        row = cursor.fetchone()
        conn.close()
        
        if row:
            try:
                return {
                    'setting_key': row['setting_key'],
                    'setting_value': json.loads(row['setting_value']),
                    'updated_at': row['updated_at']
                }
            except json.JSONDecodeError:
                return None
        return None

    def get_all_settings(self) -> Dict[str, Dict]:
        """获取所有设置"""
        conn = self.get_connection()
        cursor = conn.cursor()
        cursor.execute("SELECT setting_key, setting_value, updated_at FROM settings")
        rows = cursor.fetchall()
        conn.close()
        
        settings = {}
        for row in rows:
            try:
                settings[row['setting_key']] = json.loads(row['setting_value'])
            except json.JSONDecodeError:
                continue
        return settings

    def save_setting(self, setting_key: str, setting_value: Dict) -> bool:
        """保存设置（如果存在则更新，不存在则插入）"""
        conn = self.get_connection()
        cursor = conn.cursor()
        try:
            setting_value_json = json.dumps(setting_value, ensure_ascii=False)
            cursor.execute("""
                INSERT INTO settings (setting_key, setting_value, updated_at)
                VALUES (?, ?, ?)
                ON CONFLICT(setting_key) DO UPDATE SET
                    setting_value = excluded.setting_value,
                    updated_at = excluded.updated_at
            """, (setting_key, setting_value_json, datetime.now()))
            conn.commit()
            conn.close()
            return True
        except Exception as e:
            print(f"保存设置失败: {e}")
            conn.close()
            return False

    def save_all_settings(self, settings: Dict[str, Dict]) -> bool:
        """批量保存设置"""
        conn = self.get_connection()
        cursor = conn.cursor()
        try:
            for setting_key, setting_value in settings.items():
                setting_value_json = json.dumps(setting_value, ensure_ascii=False)
                cursor.execute("""
                    INSERT INTO settings (setting_key, setting_value, updated_at)
                    VALUES (?, ?, ?)
                    ON CONFLICT(setting_key) DO UPDATE SET
                        setting_value = excluded.setting_value,
                        updated_at = excluded.updated_at
                """, (setting_key, setting_value_json, datetime.now()))
            conn.commit()
            conn.close()
            return True
        except Exception as e:
            print(f"批量保存设置失败: {e}")
            conn.close()
            return False

    # ============ 内部辅助方法 ============
    def _ensure_settings_table(self, conn):
        """确保设置表存在"""
        cursor = conn.cursor()
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS settings (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                setting_key TEXT UNIQUE NOT NULL,
                setting_value TEXT NOT NULL,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)

    def _default_settings(self) -> Dict[str, Dict]:
        """默认设置"""
        return {
            'chat': {
                'model': 'glm-4.5-flash',
                'temperature': 0.7,
                'max_tokens': 65536
            },
            'image': {
                'model': 'cogview-3-flash',
                'size': '1024x1024'
            },
            'video': {
                'model': 'cogvideox-flash',
                'quality': 'quality',
                'with_audio': True,
                'size': '1920x1080',
                'fps': 30
            },
            'speech': {
                'model': 'glm-asr',
                'temperature': 0.95,
                'stream': False
            },
            'tts': {
                'model': 'cogtts',
                'voice': 'tongtong',
                'response_format': 'wav'
            }
        }

    def _ensure_default_settings(self):
        """确保数据库存在默认设置"""
        conn = self.get_connection()
        cursor = conn.cursor()
        try:
            cursor.execute("SELECT COUNT(*) FROM settings")
            count = cursor.fetchone()[0]
            if count == 0:
                defaults = self._default_settings()
                conn.close()
                self.save_all_settings(defaults)
            else:
                conn.close()
        except Exception:
            conn.close()
            # 如果查询失败，尝试重新初始化表并写入默认
            conn = self.get_connection()
            self._ensure_settings_table(conn)
            conn.commit()
            conn.close()
            self.save_all_settings(self._default_settings())


# 全局数据库实例
db = Database()

